package hw2.sign;

import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.Signature;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

public class DigitalSignatures {
	
	static PrivateKey sign;
	static PublicKey ver;

	/**
	 * @param args
	 */
	public static void main(String[] args) throws Exception {
		String filename = "";
		String op = "";
		switch (args.length)
		{
		case 1:
			if (args[0].equalsIgnoreCase("gen")){
				generateKeys();
				System.out.println("Keys generated.");
				System.exit(0);
			}
			else{
				filename = args[0];
			}
			break;
		case 2:
			op = args[0];
			filename = args[1];
			break;
		default:
			System.out.println("Command Line arguments are not valid.  Expected any of the following:\n"
					+ "gen\n"
					+ "sign <filename>\n"
					+ "ver <filename>\n");
			System.exit(0);
			break;
		}
		
		loadKeys();
		if (op.equalsIgnoreCase("sign")){
			sign(filename);
		}
		if (op.equalsIgnoreCase("ver")){
			if (verify(filename)){
				System.out.println("" + filename + " signature has been verified.");
			}
			else{
				System.out.println("The signature for " + filename + " is not valid.");
			}
		}
	}
	
	/**
	 * Generate a new key pair
	 */
	public static void generateKeys(){
		try{
			KeyPairGenerator gen = KeyPairGenerator.getInstance("RSA");
			gen.initialize(1024);
			KeyPair pair = gen.generateKeyPair();
			FileIO.writeBytesToFile(pair.getPrivate().getEncoded(),"rsa_sign.bin");
			FileIO.writeBytesToFile(pair.getPublic().getEncoded(),"rsa_ver.bin");
		}catch(Exception e){
			e.printStackTrace();
		}
	}
	
	/**
	 * Load keys from file system
	 */
	public static void loadKeys(){
		try{
			byte[] b = FileIO.getBytesFromFile("rsa_sign.bin");
			PKCS8EncodedKeySpec spec = new PKCS8EncodedKeySpec(b);
			sign =  KeyFactory.getInstance( "RSA" ).generatePrivate( spec );
			
			b = FileIO.getBytesFromFile("rsa_ver.bin");
			X509EncodedKeySpec spec2 = new X509EncodedKeySpec(b);
			ver =  KeyFactory.getInstance( "RSA" ).generatePublic( spec2 );
			
		}catch(Exception e){
			e.printStackTrace();
		}
	}
	
	/**
	 * Sign the file and save a signature named filename.sig
	 * @param filename
	 * @throws NoSuchAlgorithmException 
	 */
	public static void sign(String filename) throws Exception{
		Signature signer = Signature.getInstance("SHA256WITHRSA");
		signer.initSign(sign);
		byte[] tosign = FileIO.getBytesFromFile(filename);
		signer.update(tosign);
		byte[] sig = signer.sign();
		FileIO.writeBytesToFile(sig, filename + ".sig");
		System.out.println("Signed " + filename + " saved as " + filename +".sig");
	}
	
	/**
	 * signature is the file name of a *.sig file created with the sign method
	 * Returns true if the signature is valid and false if it is not valid.
	 * @param signature
	 * @return
	 */
	public static boolean verify(String filename) throws Exception{
		Signature verifier = Signature.getInstance("SHA256WITHRSA");
		verifier.initVerify(ver);
		byte[] tover = FileIO.getBytesFromFile(filename);
		verifier.update(tover);
		byte[] sig = FileIO.getBytesFromFile(filename + ".sig");
		return verifier.verify(sig);
	}
	
	

}
